"""
https://leetcode.cn/problems/managers-with-at-least-5-direct-reports/description/?envType=study-plan-v2&envId=30-days-of-pandas&lang=pythondata

570. 至少有5名直接下属的经理
已解答
中等
相关标签
premium lock icon
相关企业
提示
SQL Schema
Pandas Schema
表: Employee

+-------------+---------+
| Column Name | Type    |
+-------------+---------+
| id          | int     |
| name        | varchar |
| department  | varchar |
| managerId   | int     |
+-------------+---------+
id 是此表的主键（具有唯一值的列）。
该表的每一行表示雇员的名字、他们的部门和他们的经理的id。
如果managerId为空，则该员工没有经理。
没有员工会成为自己的管理者。
 

编写一个解决方案，找出至少有五个直接下属的经理。

以 任意顺序 返回结果表。

查询结果格式如下所示。

 

示例 1:

输入: 
Employee 表:
+-----+-------+------------+-----------+
| id  | name  | department | managerId |
+-----+-------+------------+-----------+
| 101 | John  | A          | Null      |
| 102 | Dan   | A          | 101       |
| 103 | James | A          | 101       |
| 104 | Amy   | A          | 101       |
| 105 | Anne  | A          | 101       |
| 106 | Ron   | B          | 101       |
+-----+-------+------------+-----------+
输出: 
+------+
| name |
+------+
| John |
+------+

"""

import pandas as pd

def find_managers(employee: pd.DataFrame) -> pd.DataFrame:
    managers= employee.loc[~(employee['managerId'].isnull())].groupby('managerId').size().reset_index(name='count')
    managers=managers.loc[managers['count']>=5]
    return pd.merge(employee,managers,how='inner',left_on='id',right_on='managerId').loc[:,['name']]


if __name__== "__main__":
    employee = pd.DataFrame({
        'id': [101, 102, 103, 104, 105, 106],
        'name': ['John', 'Dan', 'James', 'Amy', 'Anne', 'Ron'],
        'department': ['A', 'A', 'A', 'A', 'A', 'B'],
        'managerId': [None, 101, 101, 101, 101, 101]
    })
    print(find_managers(employee))